#!/usr/bin/python3
"""Training and Validation On Segmentation Task."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import math
import random
import shutil
import argparse
import importlib
import data_utils
import numpy as np
import pointfly as pf
import tensorflow as tf
from datetime import datetime
from data_provider import data_set_proveider
from metrics_exp import compute_mean_iou
from tf_pts_generator import pts_fts_reader, visualization
from tensorflow.python import debug as tfdbg


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--filelist', '-t', help='Path to training set ground truth (.txt)', required=True)
    parser.add_argument('--filelist_val', '-v', help='Path to validation set ground truth (.txt)', required=True)
    parser.add_argument('--load_ckpt', '-l', help='Path to a check point file for load')
    parser.add_argument('--save_folder', '-s', help='Path to folder for saving check points and summary', required=True)
    parser.add_argument('--model', '-m', help='Model to use', required=True)
    parser.add_argument('--setting', '-x', help='Setting to use', required=True)
    args = parser.parse_args()

    time_string = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
    root_folder = os.path.join(args.save_folder, '%s_%s_%s_%d' % (args.model, args.setting, time_string, os.getpid()))
    if not os.path.exists(root_folder):
        os.makedirs(root_folder)

    # sys.stdout = open(os.path.join(root_folder, 'log.txt'), 'w')

    print('PID:', os.getpid())

    print(args)

    model = importlib.import_module(args.model)
    setting_path = os.path.join(os.path.dirname(__file__), args.model)
    sys.path.append(setting_path)
    setting = importlib.import_module(args.setting)

    num_epochs = setting.num_epochs  # 1024
    batch_size = setting.batch_size  # 16
    # sample_num = setting.sample_num  # 2048
    step_val = setting.step_val  # 500
    label_weights_list = setting.label_weights  # 50 capacity contain 1.0
    rotation_range = setting.rotation_range  # [0,0,0 ,u]
    rotation_range_val = setting.rotation_range_val
    scaling_range = setting.scaling_range
    scaling_range_val = setting.scaling_range_val
    jitter = setting.jitter  # 0.001
    jitter_val = setting.jitter_val  # 0

    ######################################################################
    # Placeholders

    indices = tf.placeholder(tf.int32, shape=(None, None, 2), name="indices")
    xforms = tf.placeholder(tf.float32, shape=(None, 3, 3), name="xforms")
    rotations = tf.placeholder(tf.float32, shape=(None, 3, 3), name="rotations")
    jitter_range = tf.placeholder(tf.float32, shape=(1), name="jitter_range")
    global_step = tf.Variable(0, trainable=False, name='global_step')
    is_training = tf.placeholder(tf.bool, name='is_training')

    pts_fts = tf.placeholder(tf.float32, shape=(None, setting.point_num, setting.data_dim), name='pts_fts')
    labels_seg = tf.placeholder(tf.int64, shape=(None, setting.point_num), name='labels_seg')
    labels_weights = tf.placeholder(tf.float32, shape=(None, setting.point_num), name='labels_weights')

    ######################################################################
    # pts_fts_sampled = tf.gather_nd(pts_fts, indices=indices, name='pts_fts_sampled')
    pts_fts_sampled = pts_fts
    if setting.verbose:
        pts_fts_sampled = tf.Print(pts_fts_sampled, ["pts_fts_sampled:", tf.shape(pts_fts_sampled)])
    features_augmented = None
    if setting.data_dim > 3:
        points_sampled, features_sampled = tf.split(pts_fts_sampled,
                                                    [3, setting.data_dim - 3],
                                                    axis=-1,
                                                    name='split_points_features')
        if setting.verbose:
            features_sampled = tf.Print(features_sampled, ["features_sampled:", tf.shape(features_sampled)])
        if setting.use_extra_features:
            if setting.with_normal_feature:
                if setting.data_dim < 6:
                    print('Only 3D normals are supported!')
                    exit()
                elif setting.data_dim == 6:
                    features_augmented = pf.augment(features_sampled, rotations)
                else:
                    normals, rest = tf.split(features_sampled, [3, setting.data_dim - 6])
                    normals_augmented = pf.augment(normals, rotations)
                    features_augmented = tf.concat([normals_augmented, rest], axis=-1)
            else:
                features_augmented = features_sampled
    else:
        points_sampled = pts_fts_sampled
    points_augmented = pf.augment(points_sampled, xforms, jitter_range)

    # labels_sampled = tf.gather_nd(labels_seg, indices=indices, name='labels_sampled')
    labels_sampled = labels_seg
    # labels_weights_sampled = tf.gather_nd(labels_weights, indices=indices, name='labels_weight_sampled')
    labels_weights_sampled = labels_weights

    net = model.Net(points_augmented, features_augmented, is_training, setting)
    logits = net.logits  # (?, ?, 50)
    probs = tf.nn.softmax(logits, name='probs')
    predictions = tf.argmax(probs, axis=-1, name='predictions')
    if setting.verbose:
        predictions = tf.Print(predictions, ["predictions:", tf.shape(predictions)])

    loss_op = tf.losses.sparse_softmax_cross_entropy(labels=labels_sampled, logits=logits,
                                                     weights=labels_weights_sampled)

    with tf.name_scope('metrics'):
        loss_mean_op, loss_mean_update_op = tf.metrics.mean(loss_op)
        t_1_acc_op, t_1_acc_update_op = tf.metrics.accuracy(labels_sampled, predictions, weights=labels_weights_sampled)
        t_1_per_class_acc_op, t_1_per_class_acc_update_op = \
            tf.metrics.mean_per_class_accuracy(labels_sampled, predictions, setting.num_class,
                                               weights=labels_weights_sampled)
        ious_op, ious_up_op = tf.metrics.mean_iou(labels_sampled, predictions, setting.num_class,
                                                  weights=labels_weights_sampled)
    reset_metrics_op = tf.variables_initializer([var for var in tf.local_variables()
                                                 if var.name.split('/')[0] == 'metrics'])

    _ = tf.summary.scalar('loss/train', tensor=loss_mean_op, collections=['train'])
    _ = tf.summary.scalar('t_1_acc/train', tensor=t_1_acc_op, collections=['train'])
    _ = tf.summary.scalar('t_1_per_class_acc/train', tensor=t_1_per_class_acc_op, collections=['train'])

    _ = tf.summary.scalar('loss/val', tensor=loss_mean_op, collections=['val'])
    _ = tf.summary.scalar('t_1_acc/val', tensor=t_1_acc_op, collections=['val'])
    _ = tf.summary.scalar('t_1_per_class_acc/val', tensor=t_1_per_class_acc_op, collections=['val'])
    _ = tf.summary.scalar('t_1_per_iou/val', tensor=ious_op, collections=['val'])

    lr_exp_op = tf.train.exponential_decay(setting.learning_rate_base, global_step, setting.decay_steps,
                                           setting.decay_rate, staircase=True)
    lr_clip_op = tf.maximum(lr_exp_op, setting.learning_rate_min)
    _ = tf.summary.scalar('learning_rate', tensor=lr_clip_op, collections=['train'])
    reg_loss = setting.weight_decay * tf.losses.get_regularization_loss()
    if setting.optimizer == 'adam':
        optimizer = tf.train.AdamOptimizer(learning_rate=lr_clip_op, epsilon=setting.epsilon)
    elif setting.optimizer == 'momentum':
        optimizer = tf.train.MomentumOptimizer(learning_rate=lr_clip_op, momentum=setting.momentum, use_nesterov=True)
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = optimizer.minimize(loss_op + reg_loss, global_step=global_step)

    saver = tf.train.Saver(max_to_keep=None)

    # backup all code
    # code_folder = os.path.abspath(os.path.dirname(__file__))
    # shutil.copytree(code_folder, os.path.join(root_folder, os.path.basename(code_folder)))

    folder_ckpt = os.path.join(root_folder, 'ckpts')
    if not os.path.exists(folder_ckpt):
        os.makedirs(folder_ckpt)

    folder_summary = os.path.join(root_folder, 'summary')
    if not os.path.exists(folder_summary):
        os.makedirs(folder_summary)

    parameter_num = np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()])
    print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num))

    trainset_reader = pts_fts_reader(args.filelist)
    print("train files:{}".format(len(trainset_reader.get_tf_files())))
    filename_queue_train = tf.train.string_input_producer(trainset_reader.get_tf_files(),
                                                          num_epochs=setting.batch_num)
    pts_fts_train_op, data_num_train_op, label_train_op = pts_fts_reader.read(
        filename_queue_train, batch_size=setting.batch_size, shuffle_batch=False)

    valset_reader = pts_fts_reader(args.filelist_val)
    print("val files:{}".format(len(valset_reader.get_tf_files())))
    filename_queue_val = tf.train.string_input_producer(valset_reader.get_tf_files(),
                                                        num_epochs=setting.batch_num)
    pts_fts_val_op, data_num_val_op, label_val_op = pts_fts_reader.read(
        filename_queue_val, batch_size=setting.batch_size, shuffle_batch=False)

    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    with tf.Session() as sess:
        # sess = tfdbg.TensorBoardDebugWrapperSession(sess,"127.0.0.1:8034")
        summaries_op = tf.summary.merge_all('train')
        summaries_val_op = tf.summary.merge_all('val')
        summary_writer = tf.summary.FileWriter(folder_summary, sess.graph)

        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        # Load the model
        if args.load_ckpt is not None:
            saver.restore(sess, args.load_ckpt)
            print('{}-Checkpoint loaded from {}!'.format(datetime.now(), args.load_ckpt))

        for batch_idx_train in range(setting.batch_num):
            data_train, data_num_train, label_train = sess.run([pts_fts_train_op, data_num_train_op, label_train_op])
            if (batch_idx_train % step_val == 0 and (batch_idx_train != 0 or args.load_ckpt is not None)) \
                    or batch_idx_train == setting.batch_num - 1:
                ######################################################################
                # Validation
                filename_ckpt = os.path.join(folder_ckpt, 'iter')
                saver.save(sess, filename_ckpt, global_step=global_step)
                print('{}-Checkpoint saved to {}!'.format(datetime.now(), filename_ckpt))

                sess.run(reset_metrics_op)
                for batch_val_idx in range(setting.batch_num_val):
                    data_val, data_num_val, label_val = sess.run([pts_fts_val_op, data_num_val_op, label_val_op])
                    batch_size_val = setting.batch_size
                    points_batch = data_val
                    points_num_batch = data_num_val
                    labels_batch = label_val
                    weights_batch = np.array(label_weights_list)[labels_batch]

                    xforms_np, rotations_np = pf.get_xforms(batch_size_val,
                                                            rotation_range=rotation_range_val,
                                                            scaling_range=scaling_range_val,
                                                            order=setting.rotation_order)
                    _, _, per_class_accs, _ = sess.run(
                        [loss_mean_update_op, t_1_acc_update_op, t_1_per_class_acc_update_op, ious_up_op],
                        feed_dict={
                            pts_fts: points_batch,
                            indices: pf.get_indices(batch_size_val, sample_num, points_num_batch),
                            xforms: xforms_np,
                            rotations: rotations_np,
                            jitter_range: np.array([jitter_val]),
                            labels_seg: labels_batch,
                            labels_weights: weights_batch,
                            is_training: False,
                        })
                    sys.stdout.write(
                        '\r[ background:{:.4f} cyclist:{:.4f} tricycle:{:.4f} smallMot:{:.4f} bigMot:{:.4f} pedestrian:{:.4f} crowds:{:.4f} unknown:{:.4f} ] '.format(
                            float(per_class_accs[0]), float(per_class_accs[1]), float(per_class_accs[2]),
                            float(per_class_accs[3]),
                            float(per_class_accs[4]), float(per_class_accs[5]), float(per_class_accs[6]),
                            float(per_class_accs[7])))
                    sys.stdout.flush()
                loss_val, t_1_acc_val, t_1_per_class_acc_val, ious, summaries_val = sess.run(
                    [loss_mean_op, t_1_acc_op, t_1_per_class_acc_op, ious_op, summaries_val_op])
                summary_writer.add_summary(summaries_val, batch_idx_train)
                print('{}-[Val  ]-Average:      Loss: {:.4f}  T-1 Acc: {:.4f}  T-1 mAcc: {:.4f} ious: {:.4f}'
                      .format(datetime.now(), loss_val, t_1_acc_val, t_1_per_class_acc_val, ious))
                sys.stdout.flush()
                ######################################################################

            ######################################################################
            # Training
            points_batch = data_train
            if setting.verbose:
                print("train:{} train_num:{} train label:{}".format(data_train.shape, data_num_train.shape, label_train.shape))
                for i in range(setting.batch_size):
                    visualization(data_train[i], label_train[i],i)
            points_num_batch = data_num_train
            labels_batch = label_train
            if setting.verbose:
                print("labels_batch shape:{}".format(labels_batch.shape))
            weights_batch = np.array(label_weights_list)[labels_batch]

            offset = int(random.gauss(0, sample_num * setting.sample_num_variance))
            offset = max(offset, -sample_num * setting.sample_num_clip)
            offset = min(offset, sample_num * setting.sample_num_clip)
            sample_num_train = sample_num + offset
            xforms_np, rotations_np = pf.get_xforms(setting.batch_size,
                                                    rotation_range=rotation_range,
                                                    scaling_range=scaling_range,
                                                    order=setting.rotation_order)
            sam_index = pf.get_indices(setting.batch_size, sample_num_train, points_num_batch)
            if setting.verbose:
                print("sam_index shape:{}".format(sam_index.shape))
            sess.run(reset_metrics_op)
            _, _, _, per_class_accs = sess.run(
                [train_op, loss_mean_update_op, t_1_acc_update_op, t_1_per_class_acc_update_op],
                feed_dict={
                    pts_fts: points_batch,
                    indices: sam_index,
                    xforms: xforms_np,
                    rotations: rotations_np,
                    jitter_range: np.array([jitter]),
                    labels_seg: labels_batch,
                    labels_weights: weights_batch,
                    is_training: True,
                })
            sys.stdout.write(
                '\r[ background:{:.4f} cyclist:{:.4f} tricycle:{:.4f} smallMot:{:.4f} bigMot:{:.4f} pedestrian:{:.4f} crowds:{:.4f} unknown:{:.4f} ] '.format(
                    float(per_class_accs[0]), float(per_class_accs[1]), float(per_class_accs[2]),
                    float(per_class_accs[3]),
                    float(per_class_accs[4]), float(per_class_accs[5]), float(per_class_accs[6]),
                    float(per_class_accs[7])))
            sys.stdout.flush()
            # print("ious per class:{}".format(ious))
            # mean_iou = 0
            # for iou in ious:
            #     mean_iou +=iou
            # print("mean_iou:{}".format(mean_iou))
            if batch_idx_train % 10 == 0:
                loss, t_1_acc, t_1_per_class_acc, summaries = sess.run([loss_mean_op,
                                                                        t_1_acc_op,
                                                                        t_1_per_class_acc_op,
                                                                        summaries_op])
                summary_writer.add_summary(summaries, batch_idx_train)
                print('{}-[Train]-Iter: {:06d}  Loss: {:.4f}  T-1 Acc: {:.4f}  T-1 mAcc: {:.4f}'
                      .format(datetime.now(), batch_idx_train, loss, t_1_acc, t_1_per_class_acc))
                sys.stdout.flush()
            ######################################################################
        coord.request_stop()
        coord.join(threads)
        print('{}-Done!'.format(datetime.now()))


if __name__ == '__main__':
    main()
